package com.hujiang.basic.search.common.component;

import com.alibaba.fastjson.JSONObject;
import com.google.common.base.Preconditions;
import com.hujiang.basic.framework.core.exception.AppException;
import com.hujiang.basic.search.common.exception.EsIndexException;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.action.admin.indices.delete.DeleteIndexResponse;
import org.elasticsearch.action.admin.indices.flush.FlushRequest;
import org.elasticsearch.action.admin.indices.refresh.RefreshRequest;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkRequestBuilder;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteRequestBuilder;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexRequestBuilder;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.suggest.SuggestRequest;
import org.elasticsearch.action.suggest.SuggestResponse;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.action.update.UpdateResponse;
import org.elasticsearch.client.Client;
import org.elasticsearch.cluster.metadata.IndexMetaData;
import org.elasticsearch.common.collect.ImmutableOpenMap;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentFactory;
import org.elasticsearch.index.reindex.BulkIndexByScrollResponse;
import org.elasticsearch.index.reindex.UpdateByQueryAction;
import org.elasticsearch.index.reindex.UpdateByQueryRequest;

import java.io.IOException;
import java.math.BigDecimal;
import java.util.Map.Entry;

import static org.elasticsearch.index.query.QueryBuilders.matchAllQuery;

@Slf4j
public class HJElasticSearchClient {

    private Client client;

    public HJElasticSearchClient(Client client) {
        this.client = client;
    }

    public <T> T execute(ElasticSearchCallBack<T> action) {

        Preconditions.checkNotNull(action, "Callback object must not be null");
        try {
            return action.doInElasticSearch(client);
        } catch (Throwable e) {
            throw new AppException(e.getMessage(), e, "");
        }
    }

    public SearchResponse search(SearchRequest request) {
        log.info(JSONObject.toJSONString(request));
        return execute((c) -> c.search(request).actionGet());
    }

    public GetResponse get(GetRequest request) {
        log.info(JSONObject.toJSONString(request));
        return execute((c) -> c.get(request).actionGet());
    }

    public SearchResponse getCount(String index) {
        return execute((c) -> c.prepareSearch(index).setQuery(matchAllQuery()).setSize(0).get());
    }

    public void flushAndRefresh(String index) {
        execute((c) -> c.admin().indices().flush(new FlushRequest(index).force(true).waitIfOngoing(true)).actionGet());
        execute((c) -> c.admin().indices().refresh(new RefreshRequest(index)).actionGet());
    }


    public IndexResponse index(IndexRequest req) {
        log.info(JSONObject.toJSONString(req));
        return execute((c) -> c.index(req).actionGet());
    }

    public ImmutableOpenMap<String, IndexMetaData>  getIndices() {
        return execute((c) ->
                c.admin().cluster().prepareState().execute().actionGet().getState().getMetaData().getIndices());
    }

    public DeleteIndexResponse deleteIndex(String index) {
        return execute((c) -> c.admin().indices().prepareDelete(index).execute().actionGet());
    }

    public DeleteResponse delete(DeleteRequest req) {
        log.info(JSONObject.toJSONString(req));
        return execute((c) -> c.delete(req).actionGet());
    }

    public UpdateResponse update(UpdateRequest req) {
        log.info(JSONObject.toJSONString(req));
        return execute((c) -> c.update(req).actionGet());
    }

    public BulkResponse bulk(BulkRequestBuilder req) {
        return execute((c) -> req.execute().actionGet());
    }

    public BulkResponse bulk(BulkRequest req) {
        return execute((c) -> c.bulk(req).actionGet());
    }

    public SuggestResponse suggest(SuggestRequest req) {
        return execute((c) -> c.suggest(req).actionGet());
    }

    public BulkIndexByScrollResponse updateByQuery(UpdateByQueryRequest req) {
        log.info(JSONObject.toJSONString(req));
        return execute((c) -> c.execute(UpdateByQueryAction.INSTANCE, req).actionGet());
    }

    public BulkRequestBuilder prepareBulk() {
        return client.prepareBulk();
    }

    public IndexRequestBuilder prepareIndex(String index, String type, String id, JSONObject source) {
        try {
            XContentBuilder xb = XContentFactory.jsonBuilder().startObject();
            for (Entry<String, Object> entry : source.entrySet()) {
                xb.field(entry.getKey(), entry.getValue().getClass().isAssignableFrom(BigDecimal.class) ?
                        Double.parseDouble(entry.getValue().toString()) : entry.getValue());
            }
            xb.endObject();
            return client.prepareIndex(index, type, id).setSource(xb);
        } catch (IOException e) {
            log.error("prepareIndex error:{}",e);
            throw new EsIndexException(e.getMessage());
        }
    }

    public DeleteRequestBuilder preparedelete(String index, String type, String id) {
        return client.prepareDelete(index, type, id);
    }
}
